#include <rtc_base/byte_order.h>
#include <rtc_base/crc32.h>
#include <rtc_base/message_digest.h>
#include <rtc_base/logging.h>
#include <arpa/inet.h>

#include "ice/stun.h"
namespace xrtc
{

    const char EMPTY_TRANSACION_ID[] = "000000000000";
    const size_t STUN_FINGERPRINT_XOR_VALUE = 0x5354554e;
    const char STUN_ERROR_REASON_BAD_REQUEST[] = "Bad request";
    const char STUN_ERROR_REASON_UNAUTHORIZED[] = "Unauthorized";
    const char STUN_ERROR_REASON_SERVER_ERROR[] = "Server error";

    std::string StunMethodToString(int type)
    {

        switch (type)
        {
        case STUN_BINDING_REQUEST:
            return "BINDING REQUEST";
        case STUN_BINDING_RESPONSE:
            return "BINDING RESPONSE";
        case STUN_BINDING_ERROR_RESPONSE:
            return "BINDING ERROR_RESPONSE";
        default:
            return "Unknown<" + std::to_string(type) + ">";
        }
    }

    StunMessage::StunMessage() : type_(0), length_(0), transcation_id_(EMPTY_TRANSACION_ID)
    {
    }
    StunMessage::~StunMessage() = default;

    bool StunMessage::ValidateFingerprint(const char *data, size_t len)
    {
        // 检测长度是否是4的整数倍

        size_t fingerprint_attr_size = k_stun_attribute_header_size + StunUInt32Attribute::SIZE;
        if (len % 4 != 0 || len < k_stun_attribute_header_size + fingerprint_attr_size)
        {
            return false;
        }

        // 检查magic
        const char *magic_cookie = data + k_stun_transaction_id_offset - k_stun_magic_cookie_length;
        if (rtc::GetBE32(magic_cookie) != k_stun_magic_cookie)
        {
            return false;
        }
        // 检查 attribute 的type和length
        const char *finger_attr_data = data + len - fingerprint_attr_size;
        if (rtc::GetBE16(finger_attr_data) != STUN_ATTR_FINGERPRINT || rtc::GetBE16(finger_attr_data + sizeof(uint16_t)) != StunUInt32Attribute::SIZE)
        {
            return false;
        }

        // 提取data数据中fingerprint 的值

        uint32_t fingerprint = rtc::GetBE32(finger_attr_data + k_stun_attribute_header_size);
        return (fingerprint ^ STUN_FINGERPRINT_XOR_VALUE) == rtc::ComputeCrc32(data, len - fingerprint_attr_size);
    }

    const StunByteStringAttribute *StunMessage::GeByteString(uint16_t type)
    {
        return static_cast<const StunByteStringAttribute *>(GetAttribute(type));
    }

    const StunUInt32Attribute *StunMessage::GetUint32(uint16_t type)
    {
        return static_cast<const StunUInt32Attribute *>(GetAttribute(type));
    }

    const StunAttribute *StunMessage::GetAttribute(uint16_t type)
    {

        for (const auto &attr : attrs_)
        {

            if (attr->Type() == type)
            {
                return attr.get();
            }
        }
        return nullptr;
    }
    bool StunMessage::Read(rtc::ByteBufferReader *buf)
    {

        if (!buf)
        {
            return false;
        }
        buffer_.assign(buf->Data(), buf->Length());
        if (!buf->ReadUInt16(&type_)) // 读取之后buf 后向后偏移
        {
            return false;
        }
        // 过滤 rtp/rtcp
        if (type_ & 0x8000)
        {
            return false;
        }
        if (!buf->ReadUInt16(&length_))
        {
            return false;
        }
        std::string magic_cookie;
        if (!buf->ReadString(&magic_cookie, k_stun_magic_cookie_length))
        {
            return false;
        }
        std::string transaction_id;
        if (!buf->ReadString(&transaction_id, k_stun_transaction_id_length))
        {
            return false;
        }
        uint32_t magic_cookie_int;
        memcpy(&magic_cookie_int, magic_cookie.data(), sizeof(magic_cookie_int));
        if (rtc::NetworkToHost32(magic_cookie_int) != k_stun_magic_cookie)
        {
            // 表示就版本stun
            transaction_id.insert(0, magic_cookie);
        }
        transcation_id_ = transaction_id;

        // 剩余数据长度 【ByteBufferReader 读取之后指针会偏移】剩余的就是属性的数据
        if (buf->Length() != length_)
        {
            return false;
        }
        attrs_.resize(0);
        while (buf->Length() > 0)
        {
            uint16_t attr_type;
            uint16_t attr_length; // 属性真实的长度（不是填充之后的长度）
            if (!buf->ReadUInt16(&attr_type))
            {
                return false;
            }
            if (!buf->ReadUInt16(&attr_length))
            {
                return false;
            }
            std::unique_ptr<StunAttribute> attr(CreateAttribute(attr_type, attr_length));
            if (!attr)
            {
                if (attr_length % 4 != 0)
                {
                    attr_length += (4 - attr_length % 4); // attr_length 现在是填充之后的长度
                }
                if (!buf->Consume(attr_length)) // attr_length 填充之后长度（4字节对齐的长度）
                {
                    return false;
                }
            }
            else
            {
                if (!attr->Read(buf))
                {

                    return false;
                }
                attrs_.push_back(std::move(attr));
            }
        }

        return true;
    }

    bool StunMessage::Write(rtc::ByteBufferWriter *buf) const
    {
        if (!buf)
        {
            return false;
        }
        buf->WriteUInt16(type_);
        buf->WriteUInt16(length_);
        buf->WriteUInt32(k_stun_magic_cookie);
        buf->WriteString(transcation_id_);
        for (const auto &attr : attrs_)
        {
            buf->WriteUInt16(attr->Type());
            buf->WriteUInt16(attr->Length());
            if (!attr->Write(buf))
            {
                return false;
            }
        }
        return true;
    }

    StunAttributeValueType StunMessage::GetAttributeValueType(int type)
    {

        switch (type)
        {

        case STUN_ATTR_USERNAME:
            return STUN_VALUE_BYTE_STRING;
        case STUN_ATTR_MESSAGE_INTEGRITY:
            return STUN_VALUE_BYTE_STRING;
        case STUN_ATTR_PRIORITY:
            return STUN_VALUE_UINT32;
        default:
            return STUN_VALUE_UNKNOWN;
        }
    }
    StunAttribute *StunMessage::CreateAttribute(uint16_t type, uint16_t length)
    {
        StunAttributeValueType value_type = GetAttributeValueType(type);
        if (STUN_VALUE_UNKNOWN != value_type)
        {
            return StunAttribute::Create(value_type, type, length, this);
        }
        return nullptr;
    }
    const StunErrorCodeAttribute *StunMessage::GetErrorCode()
    {
        return static_cast<const StunErrorCodeAttribute *>(GetAttribute(STUN_ATTR_ERROR_CODE));
    }

    int StunMessage::GetErrorCodeValue()
    {
        auto error_attr = GetErrorCode();
        return error_attr ? error_attr->Code() : STUN_ERROR_GLOBAL_FAIL;
    }

    bool StunMessage::ValidateMessageIntegrityOfType(uint16_t mi_attr_type, size_t mi_attr_size, const char *data, size_t size, const std::string &password)
    {
        if (size % 4 != 0 || size < k_stun_message_integrity_size)
        {
            return false;
        }

        uint16_t length = rtc::GetBE16(&data[2]);
        if (length + k_stun_header_size != size)
        {
            return false;
        }
        // 查找mi 属性位置
        size_t current_pos = k_stun_header_size;
        bool has_message_integrity = false;

        while (current_pos + k_stun_attribute_header_size < size)
        {
            uint16_t attr_type;
            uint16_t attr_length;
            attr_type = rtc::GetBE16(&data[current_pos]);
            attr_length = rtc::GetBE16(&data[current_pos + sizeof(attr_type)]);
            if (attr_type == mi_attr_type)
            {
                has_message_integrity = true;
                break;
            }
            current_pos += k_stun_attribute_header_size + attr_length;
            // 判断是否有填充字节
            if (attr_length % 4 != 0)
            {
                current_pos += (4 - (attr_length % 4));
            }
        }
        if (!has_message_integrity)
        {
            return false;
        }
        size_t mi_pos = current_pos;

        // 组建 integrity 做哈希的内容
        std::unique_ptr<char[]> temp_data(new char[mi_pos]);
        memcpy(temp_data.get(), data, mi_pos);
        // 修改内容 stunmessage头部的长度
        if (size > current_pos + k_stun_attribute_header_size + mi_attr_size)
        {
            size_t extra_pos = mi_pos + k_stun_attribute_header_size + mi_attr_size;
            size_t extra_size = size - extra_pos;
            size_t adjust_new_len = size - extra_size - k_stun_header_size;
            rtc::SetBE16(temp_data.get() + 2, adjust_new_len);
        }
        char hmac[k_stun_message_integrity_size];
        size_t ret = rtc::ComputeHmac(rtc::DIGEST_SHA_1, password.c_str(), password.length(), temp_data.get(), current_pos, hmac, sizeof(hmac));
        if (ret != k_stun_message_integrity_size)
        {
            return false;
        }
        // 比较integrity 内容和计算内容是否一致

        return memcmp(data + mi_pos + k_stun_attribute_header_size, hmac, mi_attr_size) == 0;
    }

    StunMessage::IntegrityStatus StunMessage::ValidateMessageIntegrity(const std::string &password)
    {
        password_ = password;
        if (GeByteString(STUN_ATTR_MESSAGE_INTEGRITY))
        {
            if (ValidateMessageIntegrityOfType(STUN_ATTR_MESSAGE_INTEGRITY, k_stun_message_integrity_size, buffer_.c_str(), buffer_.length(), password))
            {
                integrity_ = IntegrityStatus::kIntegrityOk;
            }
            else
            {
                integrity_ = IntegrityStatus::kIntegrityBad;
            }
        }
        else
        {
            integrity_ = IntegrityStatus::kNotIntegrity;
        }

        return integrity_;
    }
    bool StunMessage::AddMessageIntegrity(const std::string &password)
    {

        return AddMessageIntegrityOfType(STUN_ATTR_MESSAGE_INTEGRITY, k_stun_message_integrity_size, password.c_str(), password.size());
    }
    bool StunMessage::AddFingerPrint()
    {
        auto fingerprint_attr_ptr = std::make_unique<StunUInt32Attribute>(STUN_ATTR_FINGERPRINT, 0);
        auto fingerprint_attr = fingerprint_attr_ptr.get();
        AddAttribute(std::move(fingerprint_attr_ptr));
        rtc::ByteBufferWriter buf;
        if (!Write(&buf))
        {
            return false;
        }

        size_t msg_len_for_crc32 = buf.Length() - k_stun_attribute_header_size 
        - fingerprint_attr->Length();

        uint32_t c = rtc::ComputeCrc32(buf.Data(), msg_len_for_crc32);

        fingerprint_attr->SetValue(c ^ STUN_FINGERPRINT_XOR_VALUE);
        return true;
    }
    bool StunMessage::AddMessageIntegrityOfType(uint16_t attr_type, uint16_t attr_size, const char *key, size_t key_len)
    {

        auto mi_attr_ptr = std::make_unique<StunByteStringAttribute>(attr_type, std::string(attr_size, '0'));
        auto mi_attr = mi_attr_ptr.get();

        AddAttribute(std::move(mi_attr_ptr));
        rtc::ByteBufferWriter buf;
        if (!Write(&buf))
        {

            return false;
        }

        size_t msg_len_for_hmac = buf.Length() - k_stun_attribute_header_size - attr_size;
        mi_attr->Length();
        char hmac[k_stun_message_integrity_size];

        size_t ret = rtc::ComputeHmac(rtc::DIGEST_SHA_1, key, key_len, buf.Data(), msg_len_for_hmac, hmac, sizeof(hmac));
        if (ret != sizeof(hmac))
        {
            RTC_LOG(LS_WARNING) << "compute hmac error";
            return false;
        }
        mi_attr->CopyBytes(hmac, k_stun_message_integrity_size);
        password_.assign(key, key_len);
        integrity_ = IntegrityStatus::kIntegrityOk;
        return true;
    }

    void StunMessage::AddAttribute(std::unique_ptr<StunAttribute> attr)
    {
        size_t attr_len = attr->Length();
        if (attr_len % 4 != 0)
        {
            attr_len += (4 - (attr_len % 4));
        }
        length_ += attr_len + k_stun_attribute_header_size;
        attrs_.push_back(std::move(attr));
    }

    // StunAttribute 类
    StunAttribute::StunAttribute(uint16_t type, uint16_t length) : type_(type), length_(length)
    {
    }
    StunAttribute::~StunAttribute() = default;
    StunAttribute *StunAttribute::Create(StunAttributeValueType value_type, 
        uint16_t type, uint16_t length, void */*owner*/)
    {
        switch (value_type)
        {
        case STUN_VALUE_BYTE_STRING:
            return new StunByteStringAttribute(type, length);
        case STUN_VALUE_UINT32:
            return new StunUInt32Attribute(type);
        default:
            return nullptr;
        }
    }
    void StunAttribute::ConsumePadding(rtc::ByteBufferReader *buf)
    {
        int remain = Length() % 4;
        if (remain > 0)
        {
            buf->Consume(4 - remain);
        }
    }
    std::unique_ptr<StunErrorCodeAttribute> StunAttribute::CreateErrorCode()
    {

        return std::make_unique<StunErrorCodeAttribute>(STUN_ATTR_ERROR_CODE, StunErrorCodeAttribute::MIN_SIZE);
    }
    void StunAttribute::WritePadding(rtc::ByteBufferWriter *buf)
    {
        int remain = Length() % 4;
        if (remain > 0)
        {
            char zeroes[4] = {0};
            buf->WriteBytes(zeroes, 4 - remain);
        }
    }
    // StunAttribute->StunUInt32Attribute
    StunUInt32Attribute::StunUInt32Attribute(uint16_t type) : StunAttribute(type, SIZE), bits_(0)
    {
    }
    StunUInt32Attribute::StunUInt32Attribute(uint16_t type, uint32_t value) : StunAttribute(type, SIZE), bits_(value)
    {
    }

    bool StunUInt32Attribute::Read(rtc::ByteBufferReader *buf)
    {
        if (Length() != SIZE || !buf->ReadUInt32(&bits_))
        {
            return false;
        }
        return true;
    }
    bool StunUInt32Attribute::Write(rtc::ByteBufferWriter *buf)
    {
        buf->WriteUInt32(bits_);
        return true;
    }

    // StunAttribute->StunUInt64Attribute
    StunUInt64Attribute::StunUInt64Attribute(uint16_t type) : 
    StunAttribute(type, SIZE), bits_(0)
    {
    }
    StunUInt64Attribute::StunUInt64Attribute(uint16_t type, uint64_t value) : 
    StunAttribute(type, SIZE), bits_(value)
    {
    }

    bool StunUInt64Attribute::Read(rtc::ByteBufferReader *buf)
    {
        if (Length() != SIZE || !buf->ReadUInt64(&bits_))
        {
            return false;
        }
        return true;
    }
    bool StunUInt64Attribute::Write(rtc::ByteBufferWriter *buf)
    {
        buf->WriteUInt64(bits_);
        return true;
    }

    // StunAttribute->StunByteStringAttribute
    StunByteStringAttribute::StunByteStringAttribute(uint16_t type, uint16_t length) 
    : StunAttribute(type, length)
    {
    }
    StunByteStringAttribute::~StunByteStringAttribute()
    {
        if (bytes_)
        {
            delete bytes_;
            bytes_ = nullptr;
        }
    }
    void StunByteStringAttribute::SetBytes(char *bytes)
    {
        if (bytes_)
        {
            delete[] bytes_;
        }
        bytes_ = bytes;
    }
    void StunByteStringAttribute::CopyBytes(const char *bytes, size_t len)
    {
        char *new_bytes = new char[len];
        mempcpy(new_bytes, bytes, len);
        SetBytes(new_bytes);
        SetLength(len);
    }
    StunByteStringAttribute::StunByteStringAttribute(uint16_t type, const std::string &str) : StunAttribute(type, 0)
    {
        CopyBytes(str.c_str(), str.size());
    }
    bool StunByteStringAttribute::Write(rtc::ByteBufferWriter *buf)
    {
        buf->WriteBytes(bytes_, Length());
        WritePadding(buf);

        return true;
    }

    bool StunByteStringAttribute::Read(rtc::ByteBufferReader *buf)
    {

        bytes_ = new char[Length()];
        if (!buf->ReadBytes(bytes_, Length()))
        {
            return false;
        }
        ConsumePadding(buf);
        return true;
    }

    //  StunAttribute->StunXorAddressAttribute

    StunAddressAttribute::StunAddressAttribute(uint16_t type, const rtc::SocketAddress &addr) : StunAttribute(type, 0)
    {

        SetAddress(addr);
    }
    bool StunAddressAttribute::Read(rtc::ByteBufferReader */*buf*/)
    {

        return true;
    }
    bool StunAddressAttribute::Write(rtc::ByteBufferWriter *buf)
    {
        /*
            0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9
           |x x x x x x x x|     Family    |        X-Port            |
           |                X-Address (Variable)                      |

           type:0x02
           x-port:16 位 (x-port = port^(magic cookie >>16))
           X-Address
                x-address-ipv4 = address-ipv4^magic_cookie
                x-address-ipv6 = address-ipv6^(magic_cookie 拼接 transation_id)
         */

        StunAddressFamily stun_family = Family();
        if (STUN_ADDRESS_UNDEF == stun_family)
        {
            RTC_LOG(LS_WARNING) << "write address attribute error:unknown family";
            return false;
        }

        buf->WriteUInt8(0);
        buf->WriteUInt8(stun_family);
        buf->WriteUInt16(address_.port());
        switch (address_.family())
        {
        case AF_INET:
        {
            in_addr v4addr = address_.ipaddr().ipv4_address();
            buf->WriteBytes((const char *)&v4addr, sizeof(v4addr));
            break;
        }
        case AF_INET6:
        {
            in6_addr v6addr = address_.ipaddr().ipv6_address();
            buf->WriteBytes((const char *)&v6addr, sizeof(v6addr));
            break;
        }
        default:
            return false;
        }

        return true;
    }
    StunAddressFamily StunAddressAttribute::Family()
    {
        switch (address_.family())
        {
        case AF_INET:
            return STUN_ADDRESS_IPV4;
        case AF_INET6:
            return STUN_ADDRESS_IPV6;
        default:
            return STUN_ADDRESS_UNDEF;
        }
    }
    void StunAddressAttribute::SetAddress(const rtc::SocketAddress &addr)
    {
        address_ = addr;
        switch (Family())
        {
        case STUN_ADDRESS_IPV4:
            SetLength(StunAddressAttribute::SIZE_IPV4);
            break;
        case STUN_ADDRESS_IPV6:
            SetLength(StunAddressAttribute::SIZE_IPV6);

            break;
        default:
            SetLength(StunAddressAttribute::SIZE_UNDEF);
            break;
        }
    }

    //  StunAttribute->StunXorAddressAttribute

    StunXorAddressAttribute::StunXorAddressAttribute(uint16_t type, const rtc::SocketAddress &addr) : StunAddressAttribute(type, addr)
    {
    }

    bool StunXorAddressAttribute::Write(rtc::ByteBufferWriter *buf)
    {
        StunAddressFamily stun_family = Family();
        if (STUN_ADDRESS_UNDEF == stun_family)
        {

            RTC_LOG(LS_WARNING) << "write address attribute error:unknown family";
            return false;
        }
        rtc::IPAddress xored_ip = GetXoredIp();
        if (AF_UNSPEC == xored_ip.family())
        {

            return false;
        }

        buf->WriteUInt8(0);
        buf->WriteUInt8(stun_family);
        buf->WriteUInt16(address_.port() ^ (k_stun_magic_cookie >> 16));
        switch (address_.family())
        {
        case AF_INET:
        {
            in_addr v4addr = xored_ip.ipv4_address();
            buf->WriteBytes((const char *)&v4addr, sizeof(v4addr));
            break;
        }
        case AF_INET6:
        {
            in6_addr v6addr = xored_ip.ipv6_address();
            buf->WriteBytes((const char *)&v6addr, sizeof(v6addr));
            break;
        }
        default:
            return false;
        }

        return true;
    }
    rtc::IPAddress StunXorAddressAttribute::GetXoredIp()
    {
        rtc::IPAddress ip = address_.ipaddr();
        switch (address_.family())
        {

        case AF_INET:
        {
            in_addr v4addr = ip.ipv4_address();
            v4addr.s_addr = (v4addr.s_addr ^ rtc::HostToNetwork32(k_stun_magic_cookie));
            return rtc::IPAddress(v4addr);
        }

        case AF_INET6:
            break;
        default:
            break;
        }
        return rtc::IPAddress();
    }
    bool StunXorAddressAttribute::Read(rtc::ByteBufferReader */*buf*/)
    {
        return true;
    }
    // StunAttribute->StunErrorCodeAttribute

    const uint16_t StunErrorCodeAttribute::MIN_SIZE = 4;
    StunErrorCodeAttribute::StunErrorCodeAttribute(uint16_t type, uint16_t length) : StunAttribute(type, length), class_(0), number_(0)
    {
    }
    void StunErrorCodeAttribute::SetCode(int code)
    {
        class_ = code / 100;
        number_ = code % 100;
    }

    void StunErrorCodeAttribute::SetReason(const std::string &reason)
    {

        reason_ = reason;
        SetLength(MIN_SIZE + reason.size());
    }
    bool StunErrorCodeAttribute::Read(rtc::ByteBufferReader */*buf*/)
    {
        return false;
    }
    bool StunErrorCodeAttribute::Write(rtc::ByteBufferWriter *buf)
    {
        buf->WriteUInt32(class_ << 8 | number_);
        buf->WriteString(reason_);
        WritePadding(buf);
        return true;
    }
    int StunErrorCodeAttribute::Code() const
    {
        return class_ * 100 + number_;
    }

    int GetStunSuccessResponse(int req_type)
    {

        return IsStunRequestType(req_type) ? (req_type | 0x100) : -1;
    }
    int GetStunErrorResponse(int req_type)
    {
        return IsStunRequestType(req_type) ? (req_type | 0x110) : -1;
    }
    bool IsStunRequestType(int req_type)
    {
        return (req_type & k_stun_type_mask) == 0x000;
    }

}